#1. Install packages
#install.packages('readstata13')
#install.packages('lars')
#install.packages("glmnet")
#install.packages("MASS")
#install.packages("Hmisc")


#2. Import data from STATA file (.dta) and data cleaning
rm(list=ls(all=TRUE))
library(readstata13)
data<-read.dta13("Q1-Q4 data.dta")
y<-data$TotalExp
e<-data$Electricity
w<-data$weight

ebar=mean(e)
a=0.09*ebar
b=0.01*ebar

erel=e/mean(e)

CAT=-(a+b)*erel+a
TPS=-b*erel

longes=c(CAT/data$TotalExp,TPS/data$TotalExp)
longe=c(CAT,TPS)
longdec=c(data$exp_decile,data$exp_decile)
longdata=c(rep(1,length(CAT)),rep(2,length(TPS)))
long<-data.frame(longes,longdec=as.factor(longdec),longdata=as.factor(longdata))

require(ggplot2)
longdata<-factor(longdata,rev(levels(longdata)))
png(file="figure2.png",width=1200,height=800,res=200)
ggplot(data=long,aes(x=longdec, y=longe,fill=longdata,order=as.numeric(longdata)))+
  ylab("change in household welfare (dollars)")+
  xlab("expenditure decile") +
  geom_boxplot()+
  theme(axis.title.y=element_text(vjust=-10)) +
  theme(plot.margin = unit(c(0.5,0.5,0.5,-.25), "cm"))+
  #  geom_boxplot(outlier.shape=NA)+
  #  coord_cartesian(ylim = ylim1*1.05)+
  coord_flip()+
  theme(legend.justification = c(1, 1),
        legend.position = c(.15,.25))+
  theme(legend.title = element_blank())+
  scale_x_discrete(labels=c("1"="(poorest) 1","10"="(richest) 10"))+
#  scale_fill_discrete(name="",labels,
#                      c("CAT","TPS"))
  scale_fill_grey(start=0.5,end=0.8,name="",labels,
        c("CAT","TPS"))
dev.off()

taus=seq(0,4,by=0.1)
ais=matrix(0,nrow=length(y),ncol=length(taus))
target=ais
sais=rep(0,length=length(taus))
tauCAP=sais
tauTPS=sais
dy=ave(y,data$exp_decile)
dais=ais
dsais=sais
dtarget=target
dtauCAP=tauCAP
dtauTPS=tauTPS
stauTPS=dtauTPS
stauCAP=dtauCAP
w1tauCAP=tauCAP
w1tauTPS=tauTPS
w12tauCAP=tauCAP
w12tauTPS=tauTPS
w1stauTPS=dtauTPS
w1stauCAP=dtauCAP
# alternative version of (4)
w1s2tauTPS=dtauTPS
w1s2tauCAP=dtauCAP

dTPS=ave(TPS,data$exp_decile)
dCAT=ave(CAT,data$exp_decile)
aveY=ave(y,data$exp_decile)
grandY=mean(y)
for(ntau in 1:length(taus)) {
  ais[,ntau]=y^(taus[ntau])
  sais[ntau]=sum(ais[,ntau])
  target[,ntau]=-length(y)*(ais[,ntau]/sais[ntau])*b
  dais[,ntau]=dy^(taus[ntau])
  dsais[ntau]=sum(dais[,ntau])
  dtarget[,ntau]=-length(dy)*(dais[,ntau]/dsais[ntau])*b
  tauCAP[ntau]=mean(abs(CAT-target[,ntau]))
  tauTPS[ntau]=mean(abs(TPS-target[,ntau]))
  dtauCAP[ntau]=mean(abs(CAT-dtarget[,ntau]))
  dtauTPS[ntau]=mean(abs(TPS-dtarget[,ntau]))
  stauCAP[ntau]=mean(abs(CAT-dCAT))
  stauTPS[ntau]=mean(abs(TPS-dTPS))
  
  w1tauCAP[ntau]=grandY*sqrt(mean(y/grandY*((CAT-target[,ntau])/y)^2))
  w1tauTPS[ntau]=grandY*sqrt(mean(y/grandY*((TPS-target[,ntau])/y)^2))
  w12tauCAP[ntau]=sqrt(mean(grandY/y*((CAT-target[,ntau]))^2))
  w12tauTPS[ntau]=sqrt(mean(grandY/y*((TPS-target[,ntau]))^2))
  w1stauCAP[ntau]=grandY*sqrt(mean(y/grandY*(CAT/y-dCAT/aveY)^2))
  w1stauTPS[ntau]=grandY*sqrt(mean(y/grandY*(TPS/y-dTPS/aveY)^2))

  w1s2tauCAP[ntau]=sqrt(mean(grandY/y*(CAT-dCAT)^2))
  w1s2tauTPS[ntau]=sqrt(mean(grandY/y*(TPS-dTPS)^2))
  
}


longtau=rep(taus,4)
longpenalty=c(tauCAP,tauTPS,stauCAP,stauTPS)
longpol=c(1,2,1,2)%x%rep(1,length(taus))
longcalc=c(1,1,2,2)%x%rep(1,length(taus))
long<-data.frame(longtau,longpenalty,as.factor(longpol),as.factor(longcalc))
names(long)[3]<-"Policy"
names(long)[4]<-"Calculation"

png(file="figure3a.png",width=800,height=800,res=200)
ggplot(data=long,aes(x=longtau,y=longpenalty))+
  geom_line(aes(color=long$Policy,linetype=long$Calculation),size=1.5)+
  ylab("penalty (in dollars)")+
  xlab(expression("elasticity of utility with respect to consumption ("*tau*")")) +
#  scale_color_discrete(guide=FALSE)+ (next line for b/w)
  scale_color_grey(guide=FALSE, start=0, end=0.7)+
  #  scale_linetype_discrete(name="Calculation",labels,c("total equity penalty","horizontal equity"))+
  scale_linetype_manual(name="Calculation",labels,c("total equity penalty","horizontal equity\npenalty"),
                        values=c("solid","32","11"),guide=guide_legend(keywidth=1.5))+
  annotate("text",x=3.5,y=70,label="CAT")+
  annotate("text",x=3.5,y=20,label="TPS")+
  theme(legend.justification = c(1, 1),
        legend.position = c(.5,.95))+
  scale_y_continuous(limits=c(0,110),breaks=seq(0,100,20))

#  text(3,10,"TPS")
#  scale_x_log10(breaks=c(0.1,1,10,100,1000),labels=c(0.1,1,10,100,1000))+
#  scale_color_discrete(name="HHs")
dev.off()

longtau=rep(taus,4)
longpenalty=c(w12tauCAP,w12tauTPS,w1s2tauCAP,w1s2tauTPS)
longpol=c(1,2,1,2)%x%rep(1,length(taus))
longcalc=c(1,1,2,2)%x%rep(1,length(taus))
long<-data.frame(longtau,longpenalty,as.factor(longpol),as.factor(longcalc))
names(long)[3]<-"Policy"
names(long)[4]<-"Calculation"

png(file="figure3b.png",width=800,height=800,res=200)
ggplot(data=long,aes(x=longtau,y=longpenalty))+
  geom_line(aes(color=long$Policy,linetype=long$Calculation),size=1.5)+
  ylab("penalty (in dollars)")+
  xlab(expression("elasticity of utility with respect to consumption ("*tau*")")) +
#  scale_color_discrete(guide=FALSE)+ (next line for b/w)
  scale_color_grey(guide=FALSE, start=0, end=0.7)+
  #  scale_linetype_discrete(name="Calculation",labels,c("total equity penalty","horizontal equity"))+
  scale_linetype_manual(name="Calculation",labels,c("total equity penalty","horizontal equity\npenalty"),
                        values=c("solid","32","11"),guide=FALSE)+
  annotate("text",x=3.5,y=106,label="CAT")+
  annotate("text",x=3.5,y=45,label="TPS")+
  scale_y_continuous(limits=c(0,110),breaks=seq(0,100,20))
#  text(3,10,"TPS")
#  scale_x_log10(breaks=c(0.1,1,10,100,1000),labels=c(0.1,1,10,100,1000))+
#  scale_color_discrete(name="HHs")
dev.off()

target=dtarget[,match(1,taus)]
penCAP=(abs(CAT-target))
penTPS=(abs(TPS-target))
heCAP=(abs(CAT-dCAT))
heTPS=(abs(TPS-dTPS))

w1penCAP=(grandY/y)*(CAT-target)^2
w1penTPS=(grandY/y)*(TPS-target)^2
w1heCAP=(grandY/y)*(CAT-dCAT)^2
w1heTPS=(grandY/y)*(TPS-dTPS)^2

# alternative calculation of HE...

he2CAP=(abs(CAT-dCAT))
he2TPS=(abs(TPS-dTPS))
w1he2CAP=((CAT-dCAT)^2)
w1he2TPS=((TPS-dTPS)^2)
w1he3CAP=(grandY/aveY)*((CAT-dCAT)^2)
w1he3TPS=(grandY/aveY)*((TPS-dTPS)^2)


data$CAT=CAT
data$TPS=TPS
n=14
longdata=c(aveY,dtarget[,match(1,taus)],CAT,w1he2CAP,penCAP,heCAP,w1penCAP,w1heCAP,TPS,w1he2TPS,penTPS,heTPS,w1penTPS,w1heTPS)
perc=c(4,7,8,10,13,14)
longdec=rep(1,n)%x%data$exp_decile
longgr=seq(1,n)%x%rep(1,length(CAT))
long<-data.frame(longdata,longdec,longgr)

deciletable=aggregate(long$longdata,list(long$longdec,long$longgr),mean)
deciletable=matrix(deciletable[,3],10,n)
deciletable[,perc]=sqrt(deciletable[,perc])
deciletable=deciletable[seq(10,1,-1),]

agg=deciletable
agg[,perc]=agg[,perc]^2
cts=table(data$exp_decile)
agg=cts[seq(10,1-1)]%*%agg
agg=agg/sum(cts)
agg[,perc]=sqrt(agg[,perc])

deciletable=rbind(deciletable,agg)

write.csv(deciletable,file="table3.csv")
